/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.fn.harness;



import com.bff.gaia.unified.fn.harness.control.AddHarnessIdInterceptor;

import com.bff.gaia.unified.fn.harness.control.UnifiedFnControlClient;

import com.bff.gaia.unified.fn.harness.control.ProcessBundleHandler;

import com.bff.gaia.unified.fn.harness.control.RegisterHandler;

import com.bff.gaia.unified.fn.harness.data.UnifiedFnDataGrpcClient;

import com.bff.gaia.unified.fn.harness.logging.UnifiedFnLoggingClient;

import com.bff.gaia.unified.fn.harness.state.UnifiedFnStateGrpcClientCache;

import com.bff.gaia.unified.fn.harness.stream.HarnessStreamObserverFactories;

import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi;

import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi.InstructionRequest;

import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi.InstructionResponse.Builder;

import com.bff.gaia.unified.model.pipeline.v1.Endpoints;

import com.bff.gaia.unified.runners.core.construction.PipelineOptionsTranslation;

import com.bff.gaia.unified.sdk.extensions.gcp.options.GcsOptions;

import com.bff.gaia.unified.sdk.fn.IdGenerator;

import com.bff.gaia.unified.sdk.fn.IdGenerators;

import com.bff.gaia.unified.sdk.fn.channel.ManagedChannelFactory;

import com.bff.gaia.unified.sdk.fn.stream.OutboundObserverFactory;

import com.bff.gaia.unified.sdk.function.ThrowingFunction;

import com.bff.gaia.unified.sdk.io.FileSystems;

import com.bff.gaia.unified.sdk.options.ExperimentalOptions;

import com.bff.gaia.unified.sdk.options.PipelineOptions;

import com.bff.gaia.unified.vendor.grpc.v1p13p1.com.google.protobuf.TextFormat;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableList;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import java.util.EnumMap;

import java.util.List;



/**

 * Main entry point into the Unified SDK Fn Harness for Java.

 *

 * <p>This entry point expects the following environment variables:

 *

 * <ul>

 *   <li>HARNESS_ID: A String representing the ID of this FnHarness. This will be added to the

 *       headers of calls to the Unified Control Service

 *   <li>LOGGING_API_SERVICE_DESCRIPTOR: A {@link

 *       com.bff.gaia.unified.model.pipeline.v1.Endpoints.ApiServiceDescriptor} encoded as text

 *       representing the endpoint that is to be connected to for the Unified Fn Logging service.

 *   <li>CONTROL_API_SERVICE_DESCRIPTOR: A {@link Endpoints.ApiServiceDescriptor} encoded as text

 *       representing the endpoint that is to be connected to for the Unified Fn Control service.

 *   <li>PIPELINE_OPTIONS: A serialized form of {@link PipelineOptions}. See {@link PipelineOptions}

 *       for further details.

 * </ul>

 */

public class FnHarness {

  private static final String HARNESS_ID = "HARNESS_ID";

  private static final String CONTROL_API_SERVICE_DESCRIPTOR = "CONTROL_API_SERVICE_DESCRIPTOR";

  private static final String LOGGING_API_SERVICE_DESCRIPTOR = "LOGGING_API_SERVICE_DESCRIPTOR";

  private static final String PIPELINE_OPTIONS = "PIPELINE_OPTIONS";

  private static final Logger LOG = LoggerFactory.getLogger(FnHarness.class);



  private static Endpoints.ApiServiceDescriptor getApiServiceDescriptor(String env)

      throws TextFormat.ParseException {

    Endpoints.ApiServiceDescriptor.Builder apiServiceDescriptorBuilder =

        Endpoints.ApiServiceDescriptor.newBuilder();

    TextFormat.merge(System.getenv(env), apiServiceDescriptorBuilder);

    return apiServiceDescriptorBuilder.build();

  }



  public static void main(String[] args) throws Exception {

    System.out.format("SDK Fn Harness started%n");

    System.out.format("Harness ID %s%n", System.getenv(HARNESS_ID));

    System.out.format("Logging location %s%n", System.getenv(LOGGING_API_SERVICE_DESCRIPTOR));

    System.out.format("Control location %s%n", System.getenv(CONTROL_API_SERVICE_DESCRIPTOR));

    System.out.format("Pipeline options %s%n", System.getenv(PIPELINE_OPTIONS));



    String id = System.getenv(HARNESS_ID);

    PipelineOptions options = PipelineOptionsTranslation.fromJson(System.getenv(PIPELINE_OPTIONS));



    Endpoints.ApiServiceDescriptor loggingApiServiceDescriptor =

        getApiServiceDescriptor(LOGGING_API_SERVICE_DESCRIPTOR);



    Endpoints.ApiServiceDescriptor controlApiServiceDescriptor =

        getApiServiceDescriptor(CONTROL_API_SERVICE_DESCRIPTOR);



    main(id, options, loggingApiServiceDescriptor, controlApiServiceDescriptor);

  }



  public static void main(

      String id,

      PipelineOptions options,

      Endpoints.ApiServiceDescriptor loggingApiServiceDescriptor,

      Endpoints.ApiServiceDescriptor controlApiServiceDescriptor)

      throws Exception {

    ManagedChannelFactory channelFactory;

    List<String> experiments = options.as(ExperimentalOptions.class).getExperiments();

    if (experiments != null && experiments.contains("unified_fn_api_epoll")) {

      channelFactory = ManagedChannelFactory.createEpoll();

    } else {

      channelFactory = ManagedChannelFactory.createDefault();

    }

    OutboundObserverFactory outboundObserverFactory =

        HarnessStreamObserverFactories.fromOptions(options);

    channelFactory =

        channelFactory.withInterceptors(ImmutableList.of(AddHarnessIdInterceptor.create(id)));

    main(

        id,

        options,

        loggingApiServiceDescriptor,

        controlApiServiceDescriptor,

        channelFactory,

        outboundObserverFactory);

  }



  public static void main(

      String id,

      PipelineOptions options,

      Endpoints.ApiServiceDescriptor loggingApiServiceDescriptor,

      Endpoints.ApiServiceDescriptor controlApiServiceDescriptor,

      ManagedChannelFactory channelFactory,

      OutboundObserverFactory outboundObserverFactory)

      throws Exception {

    IdGenerator idGenerator = IdGenerators.decrementingLongs();

    // The logging client variable is not used per se, but during its lifetime (until close()) it

    // intercepts logging and sends it to the logging service.

    try (UnifiedFnLoggingClient logging =

        new UnifiedFnLoggingClient(

            options, loggingApiServiceDescriptor, channelFactory::forDescriptor)) {



      LOG.info("Fn Harness started");

      // Register standard file systems.

      FileSystems.setDefaultPipelineOptions(options);

      EnumMap<

              UnifiedFnApi.InstructionRequest.RequestCase,

		  ThrowingFunction<InstructionRequest, Builder>>

          handlers = new EnumMap<>(UnifiedFnApi.InstructionRequest.RequestCase.class);



      RegisterHandler fnApiRegistry = new RegisterHandler();

      UnifiedFnDataGrpcClient unifiedFnDataMultiplexer =

          new UnifiedFnDataGrpcClient(options, channelFactory::forDescriptor, outboundObserverFactory);



      UnifiedFnStateGrpcClientCache unifiedFnStateGrpcClientCache =

          new UnifiedFnStateGrpcClientCache(

              idGenerator, channelFactory::forDescriptor, outboundObserverFactory);



      ProcessBundleHandler processBundleHandler =

          new ProcessBundleHandler(

              options, fnApiRegistry::getById, unifiedFnDataMultiplexer, unifiedFnStateGrpcClientCache);

      handlers.put(UnifiedFnApi.InstructionRequest.RequestCase.REGISTER, fnApiRegistry::register);

      // TODO(BEAM-6597): Collect MonitoringInfos in ProcessBundleProgressResponses.

      handlers.put(

          UnifiedFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE,

          processBundleHandler::processBundle);

      UnifiedFnControlClient control =

          new UnifiedFnControlClient(

              id, controlApiServiceDescriptor, channelFactory, outboundObserverFactory, handlers);



      LOG.info("Entering instruction processing loop");

      control.processInstructionRequests(options.as(GcsOptions.class).getExecutorService());

    } finally {

      System.out.println("Shutting SDK harness down.");

    }

  }

}